import random
import backtrader as bt
from loguru import logger
import numpy as np
from .boll_m import CustomBollingerBands  # 新增导入
from datetime import datetime
import json
from pathlib import Path

class StockInBollmUpChannel(bt.Indicator):
    lines = ('bollm_in_upchannel',)
    params = (
        ('days_count', 6),
        ('threshold', 4),
    )
    def __init__(self):
        self.bollm=CustomBollingerBands(self.data)
    def next(self):
        score_boll = 0
        score_ub = 0
        score_lb = 0
        for i in range(self.params.days_count):
            if self.bollm.lines.boll[-i] > self.bollm.lines.boll[-i - 1]:
                score_boll += 1
            if self.bollm.lines.ub[-i] > self.bollm.lines.ub[-i - 1]:
                score_ub += 1
            if self.bollm.lines.lb[-i] > self.bollm.lines.lb[-i - 1]:
                score_lb += 1
        
        bollm_in_upchannel=( score_boll >= self.params.threshold and score_ub >= self.params.threshold and score_lb >= self.params.threshold)
        self.lines.bollm_in_upchannel[0] = bollm_in_upchannel

class UncleGenStrategy(bt.Strategy):
    params = (
        ('target_trade_date', "2025-07-25"),
        ('days_count', 6),
        ('threshold', 4),
    )
    def __init__(self):
        if type(self.params.target_trade_date) == str:
            # 支持多种日期格式
            date_str = self.params.target_trade_date
            if len(date_str) == 8 and date_str.isdigit():
                # YYYYMMDD格式
                date_str = f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:8]}"
            try:
                self.params.target_trade_date = datetime.strptime(date_str, '%Y-%m-%d').date()
            except ValueError as e:
                logger.error(f"日期格式错误: {self.params.target_trade_date}, 错误: {e}")
                # 使用默认日期
                self.params.target_trade_date = datetime.now().date()
        self.stocks = self.datas
        self.last_date = None
        self.prenext_count = 0
        self.next_count = 0

        for data in self.stocks:
            data.bollm = CustomBollingerBands(data)
            data.macd = bt.indicators.MACD(data)
            # logger.debug(f"为 {data._name} 添加布林带指标、macd指标")

        logger.info(f"策略初始化完成，共加载 {len(self.stocks)} 只股票")

    def work(self):
        dt = self.datas[0].datetime.date(0)
        if dt != self.params.target_trade_date:
            # logger.debug(f"当前日期: {dt} 不是目标日期: {self.params.target_trade_date}")
            return
        logger.debug(f"当前日期: {dt} 是目标日期: {self.params.target_trade_date}")

        result=dict()
        result['target_trade_date']=str(dt)
        result['stock_meet_all']=[]

        for data in self.stocks:
            if len(data) > 0:
                current_result=dict()
                result[data._name]=current_result
                try:
                    score_boll = 0
                    score_ub = 0
                    score_lb = 0
                    score_close = 0
                    for i in range(self.params.days_count):
                        if data.bollm.lines.boll[-i] > data.bollm.lines.boll[-i - 1]:
                            score_boll += 1
                        if data.bollm.lines.ub[-i] > data.bollm.lines.ub[-i - 1]:
                            score_ub += 1
                        if data.bollm.lines.lb[-i] > data.bollm.lines.lb[-i - 1]:
                            score_lb += 1
                        if data.close[-i] > data.close[-i - 1]:
                            score_close += 1
                    current_result['bollm_in_upchannel']=( score_boll >= self.params.threshold and score_ub >= self.params.threshold and score_lb >= self.params.threshold)
                    current_result['k_up']=(score_close >= self.params.threshold)

                    current_result['high_above_bollm']=(data.high[0] > data.bollm.lines.boll[0])

                    current_result['high_under_ub']=(data.high[0] < data.bollm.lines.ub[0]) # 条件9

                    latest_25_volumes=[data.volume[i] for i in range(-24, 1, 1)]
                    latest_25_volumes_sorted=sorted(latest_25_volumes)
                    latest_3_volumes=[data.volume[i] for i in range(-2, 1, 1)]
                    current_result['volume_in_top_5_lowest_volumes']=any(volume in latest_25_volumes_sorted[0:5] for volume in latest_3_volumes)

                    current_result['macd_positive']=(data.macd.macd[-1] > data.macd.signal[-1])

                    if all(current_result.values()):
                        result['stock_meet_all'].append(data._name)
                        logger.info(f"{data._name} 符合条件")

                except IndexError as error:
                    logger.error(f"{data._name}:{error}")
        logger.info(f"符合条件的股票数量: {len(result['stock_meet_all'])}")

        # 把result写入文件，路径是{project_root}/data/result0.json。使用pathlib
        project_root = Path(__file__).parent.parent.parent
        result_file = project_root / 'data' / 'result0.json'
        result_file.parent.mkdir(exist_ok=True)  # 确保data目录存在
        with open(result_file, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=2)



    def prenext(self):
        self.prenext_count += 1
        self.work()

    def next(self):
        self.next_count += 1
        self.work()
    
    def stop(self):
        logger.info(f"策略执行完成 - prenext调用: {self.prenext_count} 次, next调用: {self.next_count} 次") 